"""
基于 LoRA（Low-Rank Adaptation）技术对 Qwen2.5-1.5B-Instruct 模型进行法律罪名分类任务微调的完整流程。
"""
import json

from transformers import TrainingArguments
from transformers import Trainer
from transformers import AutoTokenizer, AutoConfig, AutoModelForSequenceClassification
import torch
from peft import get_peft_model, LoraConfig


# os.environ["TF_ENABLE_ONEDNN_OPTS"]="0"

# 数据批处理模块
def data_collator(features) -> dict:
	"""
	功能：将一批数据整理为模型可接受的输入格式（统一长度的 token 序列 + 标签）。
	核心逻辑：先计算批次内最长序列长度，再对短序列进行填充，确保批次内所有数据长度一致。
	:param features
	:return:
	"""
	input_list = []
	labels_list = []
	max_seq_length = -1
	
	# 第一步：确定批次内最长序列长度（过滤超1000 tokens 的长文本避免显存溢出）
	for feature in features:
		ids = feature["input_ids"]["input"]
		if len(ids) > 1000:
			continue
		max_seq_length = max(len(ids), max_seq_length)
	
	# 第二步：对所有序列进行长度补齐（用 0 前补，统一批次长度）
	for feature in features:
		ids = feature["input_ids"]["input"]
		if len(ids) > 1000:
			continue
		label = feature["input_ids"]["output"]  # 罪名对应的 id 标签
		ids = (max_seq_length - len(ids)) * [0] + ids  # 长度补齐，用 0 前补
		input_list.append(torch.LongTensor(ids))
		labels_list.append(torch.LongTensor([label]))
	
	input_ids = torch.stack(input_list)
	labels = torch.stack(labels_list)
	
	return {
			"input_ids": input_ids,
			"labels":    labels,
	}


# 加载罪名 id 映射表
with open("crime_types/accusation_id", encoding="utf-8") as f:
	# json -> dict
	accusation_id = json.load(f)

# model_name = "E:\\code\\qwen2.5-1.5B-Instruct"
model_name = "Qwen/Qwen2.5-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)  # 加载分词器
# 加载分类模型
model = AutoModelForSequenceClassification.from_pretrained(
		model_name,  # "Qwen/Qwen2.5-1.5B-Instruct"
		num_labels=len(accusation_id),  # 分类数 202 种罪名
		load_in_8bit=True,  # 使用 8-bit 量化加载, 加速推理
		trust_remote_code=True,  # 信任远程代码，避免下载失败
		device_map="auto",  # 自动分配设备
)
# 加载预训练模型的配置（config）: 只加载结构与参数设置
config = AutoConfig.from_pretrained(
		model_name,  # 预训练模型名称或本地路径
		trust_remote_code=True,  # 允许使用模型仓库自定义的配置代码（例如 Qwen 系列）
		device_map='auto'  # 指定设备分配策略（常用于大模型自动分配 GPU/CPU）
)

# LoRA 配置
lora_config = LoraConfig(
		r=16,  # LoRA 矩阵的秩（低秩分解的维度）
		lora_alpha=16,  # LoRA 缩放因子（控制更新幅度）
		target_modules=["q_proj", "v_proj"],  # 指定对哪些子层应用 LoRA（这里是注意力的Q、V）
		lora_dropout=0.1,  # LoRA 层的 dropout，增强泛化能力
		bias="none",  # 不训练原始偏置参数
)

# 训练参数配置, Hugging Face transformers 中 TrainingArguments 的标准用法，
# 用来控制训练流程的所有超参数和保存/日志策略。
training_args = TrainingArguments(
		output_dir="qwen_lora_output",  # ✅ 保存模型与检查点的目录
		learning_rate=2e-4,  # ✅ 初始学习率（LoRA 常用 1e-4 ~ 5e-4）
		per_device_train_batch_size=4,  # ✅ 每个设备（如每块GPU）的 batch 大小
		gradient_accumulation_steps=4,  # ✅ 梯度累积，等效于总 batch size = 4×4=16
		num_train_epochs=3,  # ✅ 总训练轮数
		weight_decay=0.01,  # ✅ L2 正则，防止过拟合
		logging_dir="./logs",  # ✅ TensorBoard 日志输出目录
		logging_steps=10,  # ✅ 每 10 step 打一次 log
		save_strategy="steps",  # ✅ 按 step 保存模型
		save_steps=30,  # ✅ 每 30 step 保存一次 checkpoint
		fp16=True,  # ✅ 启用混合精度（节省显存、加速）
		push_to_hub=False,  # ❌ 不推送到 Hugging Face Hub
		report_to="tensorboard",  # ✅ 日志系统（支持 tensorboard / wandb / none）
		save_total_limit=2,  # ✅ 只保留最近 2 个 checkpoint，自动清理旧的
)
# 将模型（model_weight）的填充符（pad token）ID 与分词器（tokenizer）的填充符 ID 同步
# 让模型正确识别并忽略填充的 token，避免训练或推理时把 padding 当成有效内容。
model.config.pad_token_id = tokenizer.pad_token_id

with open("data/train_data_classify_token", encoding="utf-8") as f:
	dataset = [{"input_ids": json.loads(s)} for s in f.readlines()]

# 把 lora 挂在到模型上
model = get_peft_model(model, lora_config)
# print(model_weight)

# 使用 lora 的时候，它会默认把原模型的所有参数都冻结住
# 把被误冻结的层，重新恢复
for name, param in model.score.named_parameters():
	param.requires_grad = True


training_args.report_to = None
# 把 模型 + 数据 + 配置 给它，就能一键训练
trainer = Trainer(
		model=model,
		args=training_args,
		train_dataset=dataset,
		data_collator=data_collator,
)

# 启动训练
trainer.train()

# 保存 LoRA 模型
model.save_pretrained("./qwen_lora_model")
# ./qwen_lora_model/adapter_model.bin：LoRA 权重文件（原始模型参数不会被保存）。
# ./qwen_lora_model/adapter_config.json：LoRA 配置（如 rank、alpha 等参数）。

# 保存分类层
torch.save(model.score.state_dict(), "model_weight/score_weights.pt")
